% MICRO-373 Advanced Microfabrication Practicals
% Reference implementation of the Gerchberg-Saxton algorithm
% for 3D printing of phase mask with Nanoscibe
% This is an implementation of the Gerchberg–Saxton algorithm to calculate
% the phase mask to produce an target intensity image at the Fourier plane.
% The two known conditions are 1. the intensity map of the target 2. the
% unitary amplitude of the phase mask.

% 12.02.2025 Ye Pu

% Clear the workspace

close all;
clear;

% Target image

fname = "EPFL-logo-512.png";    % File name of target image

T = double(imread(fname));      % Load target image
P = ones(size(T));              % Phase angle at Fourier domain
T = T/max(max(T));              % Normalized target intensity
meanT = mean(mean(T));          % Mean intensity of target
sumT = sum(sum(T));             % Sum intensity of target

% Iteration control

Niter = 100;                    % Number of iterations
Ndstep = 10;                    % Display results every Ndstep iterations

% Prepare for results display

figure;                         % Figure to be split into 3:
                                % #1: phase angle at phase mask plane
                                % #2: amplitude at target plane
                                % #3: error from last iteration

% Gerchberg-Saxton iteration

T = ifftshift(T);               % Pre-transform
ST = T;                         % Iteration image

for n = 1:Niter
    
    % From the target to the phase mask

    FT = fft2(ST);              % Fourier transform from target to phase mask 
    P = angle(FT);              % Retrieve the phase angle
    
    % From the phase mask to the getget

    CP = exp(1i*P);             % Enforce unitary amplitude
    IFCP = ifft2(CP);           % Inverse transform from phase mask to target
    STlast = abs(IFCP).^2;      % Record the intensity for error calculation
    s = sum(sum(STlast));       % Sum intensity this iteration
    q = sumT/s;                 % Intensity preservation factor
    STlast = STlast*q;          % Preserve sum intensity
    ST = T.*exp(1i*angle(IFCP));    % Enforce target amplitude

    % Calculate the iteration error
    
    STerror = T-STlast;         % Calculate the change
    E(n) = std2(STerror);       % Summrize the error in standard deviation

    if mod(n,Ndstep) == 0       % Display results every Ndstep iterations
        subplot(1,3,1);         % Show phase angle in panel 1
        imagesc(fftshift(P));
        axis equal;
        axis off;
        drawnow;

        subplot(1,3,2);         % show intensity map in panel 2
        imagesc(fftshift(STlast));
        axis equal;
        axis off;
        drawnow;

        subplot(1,3,3);         % show error curve in panel 3
        plot(E,'.-');
        
        fprintf("Iteration #%05d: error = %7f5\n", n, E(n));
    end
end

% Generate data (height image) file for Nanoscribe 

data = (fftshift(P)+pi)/(2*pi);

[~,datafn,~] = fileparts(fname);
datafn = "DOE_"+datafn+".png";
imwrite(data,datafn);